# -*- coding: utf-8 -*-
# @Time    : 2021/9/30 10:29
# @Author  : Jonny
# @File    : LuckDbSqliteHelperUtils.py

import sqlite3
from timeit import default_timer
import logging


# ---- 使用 with 的方式来优化代码
class LuckDbSqliteHelperUtils(object):

    def __init__(self, db_name, commit=True, log_time=True, log_label='总用时'):
        self.db_name = db_name
        """
        :param commit: 是否在最后提交事务(设置为False的时候方便单元测试)
        :param log_time:  是否打印程序运行总时间
        :param log_label:  自定义log的文字
        """
        self._log_time = log_time
        self._commit = commit
        self._log_label = log_label

    # 创建数据库连接
    def get_connection(self):
        # 创建数据库连接
        conn = sqlite3.connect(self.db_name)
        return conn

    def __enter__(self):
        # 如果需要记录时间
        if self._log_time is True:
            self._start = default_timer()

        # 在进入的时候自动获取连接和cursor
        conn = self.get_connection()
        # 创建光标对象
        cursor = conn.cursor()

        self._conn = conn
        self._cursor = cursor
        return self

    def __exit__(self, *exc_info):
        # 提交事务
        if self._commit:
            self._conn.commit()
        # 在退出的时候自动关闭连接和cursor
        self._cursor.close()
        self._conn.close()

        if self._log_time is True:
            diff = default_timer() - self._start
            logging.info(f'-- {self._log_label}: {diff} 秒')

    @property
    def cursor(self):
        return self._cursor

    def execute_query_one(self, sql, params=None):
        """
        执行查询并返回单条记录
        :param sql: SQL语句
        :param params: 参数
        :return: 单条记录或None
        """
        with LuckDbSqliteHelperUtils(self.db_name) as db:
            try:
                if params:
                    db.cursor.execute(sql, params)
                else:
                    db.cursor.execute(sql)
                return db.cursor.fetchone()
            except Exception as e:
                logging.error(f"执行查询失败: {e}")
                return None

    def execute_query_all(self, sql, params=None):
        """
        执行查询并返回所有记录
        :param sql: SQL语句
        :param params: 参数
        :return: 所有记录列表
        """
        with LuckDbSqliteHelperUtils(self.db_name) as db:
            try:
                if params:
                    db.cursor.execute(sql, params)
                else:
                    db.cursor.execute(sql)
                return db.cursor.fetchall()
            except Exception as e:
                logging.error(f"执行查询失败: {e}")
                return []

    def execute_non_query(self, sql, params=None):
        """
        执行非查询SQL（INSERT, UPDATE, DELETE）
        :param sql: SQL语句
        :param params: 参数
        :return: 影响的行数
        """
        with LuckDbSqliteHelperUtils(self.db_name) as db:
            try:
                if params:
                    result = db.cursor.execute(sql, params)
                else:
                    result = db.cursor.execute(sql)
                return db.cursor.rowcount
            except Exception as e:
                logging.error(f"执行非查询失败: {e}")
                return 0

    def execute_many(self, sql, params_list):
        """
        批量执行SQL
        :param sql: SQL语句
        :param params_list: 参数列表
        :return: 影响的行数
        """
        with LuckDbSqliteHelperUtils(self.db_name) as db:
            try:
                db.cursor.executemany(sql, params_list)
                return db.cursor.rowcount
            except Exception as e:
                logging.error(f"批量执行失败: {e}")
                return 0

    def get_last_insert_id(self):
        """
        获取最后插入记录的ID
        :return: 最后插入的ID
        """
        return self._cursor.lastrowid
    
    # 便捷方法，用于兼容旧代码
    def execute_query(self, sql, params=None):
        """执行查询并返回所有记录（兼容方法）"""
        return self.execute_query_all(sql, params)
    
    def execute_update(self, sql, params=None):
        """执行更新操作（兼容方法）"""
        return self.execute_non_query(sql, params)
